pacman::p_load(tidyverse,data.table,broom,modelsummary,kableExtra,latex2exp,ggpubr,lfe,fixest)
rm(list=ls())
options("modelsummary_format_numeric_latex" = "plain")
################################################################################

###### Data

spatial_id_a = 'county_id'  
spatial_id_b = 'amc_id'

df = fread("inputs/data_supply_within.csv.gz")
df$y = log(df$relative_acreage)
df$x = log(df$relative_payoff)
df$iv = log(df$relative_shock)
df_sample_xy = df[is.finite(x) & is.finite(y),]

df = df[is.finite(x) & is.finite(y) & is.finite(df$iv),]
df$location_fe = df$spatial_id
df$frontier=0
frontier_list = c('NORTE','NORDESTE','NOA','NEA','PATAGONIA','CUYO')
df[region %in% frontier_list,]$frontier=1

###### Regressions

m_ols = felm(y ~ x | year + location_fe | 0 | spatial_id, data=df)
m_ols_i = felm(y ~ x + x:frontier | year + location_fe | 0 | spatial_id, data=df)
m_iv = felm(y ~ 1 | year + location_fe | (x~iv) | spatial_id, data=df)
m_iv_i = felm(y ~ 1 + x:frontier | year + location_fe | (x~iv) | spatial_id, data=df)
fs_iv = as.character(round(m_iv$stage1$iv1fstat$x[5],1))
fs_iv_i = as.character(round(m_iv_i$stage1$iv1fstat$x[5],1))

###### Table
tab_title = 'Nested model - substitution elasticity between commodities (within-nest).'
rows <- tribble(~term,               ~OLS, ~IV, ~OLS, ~IV,
                'Time FE'             ,'$\\checkmark$', '$\\checkmark$', '$\\checkmark$', '$\\checkmark$',
                'Location FE'         ,'$\\checkmark$', '$\\checkmark$', '$\\checkmark$', '$\\checkmark$',
                'KP-Wald First stage F-statistic'  ,'', fs_iv, '' , fs_iv_i)
attr(rows, 'position') <- c(5,6)
gm <- tibble::tribble(
  ~raw,        ~clean,          ~fmt,
  "nobs",      "Observations",  0)
final_table_tex <- msummary(list('OLS'=m_ols,'IV'=m_iv,'OLS'=m_ols_i,'IV'=m_iv_i),
                         title = paste0('\\label{tab: estimation - supply - within}',tab_title),
                         coef_omit = 'Intercept',
                         coef_map =  c('x'   = '$\\frac{\\theta}{\\lambda}$', 'x(fit)' ='$\\frac{\\theta}{\\lambda}$', 'x:frontier'='$\\frac{\\theta}{\\lambda}$ $\\times$ frontier region'),
                         estimate = "{estimate}{stars}",
                         gof_map = gm, 
                         escape=FALSE,
                         add_rows = rows,
                         output='latex')
final_table_tex = final_table_tex %>% kable_styling(latex_options = "HOLD_position")
kableExtra::save_kable(final_table_tex, file = "../../output/tables/table3a_within.tex")

###### Export parameters
theta_lambda = m_iv_i$coefficients[2]
theta_lambda_frontier = m_iv_i$coefficients[1]
theta_lambda_mf = theta_lambda
theta_lambda_mf_frontier = theta_lambda+theta_lambda_frontier

lambda_mf_frontier_guess = 0.05
rescale_k = (theta_lambda_mf_frontier*lambda_mf_frontier_guess)^(-0.5)-1
rescale_theta_lambda = theta_lambda_mf_frontier*rescale_k
rescale_lambda = lambda_mf_frontier_guess*rescale_k
scale_spec = data.table(rbind(c(rescale_theta_lambda,rescale_lambda)))
setnames(scale_spec, new=c('rescale_theta_lambda','rescale_lambda'))
write.csv(scale_spec, 'inputs/spec_scale_within.csv', row.names = FALSE)

df = fread('../../data/landuse/clean/geographicunits_argentina.csv.gz')
setnames(df, old=spatial_id_a, new='spatial_id')
df_a = unique(df[, list(spatial_id, region)])
df = fread('../../data/landuse/clean/geographicunits_brazil.csv.gz')
setnames(df, old=spatial_id_b, new='spatial_id')
df_b = unique(df[, list(spatial_id, region)])
df = rbind(df_a,df_b)
df$theta_lambda = theta_lambda_mf
df[region %in% frontier_list,]$theta_lambda = theta_lambda_mf_frontier
df$theta_lambda_rescaled = df$theta_lambda+rescale_theta_lambda
df_within_elasticities = df
write.csv(df_within_elasticities, "inputs/parameters_supply_within_tl.csv", row.names = FALSE)

df=merge(df_sample_xy, df_within_elasticities, by=c('spatial_id','region'), all.x=T)
df$log_zetat_ratio = (1/df$theta_lambda)*df$y-df$x
df$zetat_ratio = exp(df$log_zetat_ratio)
df$log_zetat_ratio_rescaled = (1/(df$theta_lambda_rescaled))*df$y-df$x
df$zetat_ratio_rescaled = exp(df$log_zetat_ratio_rescaled)
df_zetat = df[,c("spatial_id","region","year","country","landuse","state_id",
                 "theta_lambda_rescaled","log_zetat_ratio","zetat_ratio",
                 "log_zetat_ratio_rescaled","zetat_ratio_rescaled")] 
write.csv(df_zetat, "inputs/parameters_supply_within_zetat_ratio.csv", row.names = FALSE)

